package com.open.basic.schedule;

import com.open.basic.constant.TraceId;
import org.slf4j.MDC;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;

import javax.validation.constraints.NotNull;
import java.util.UUID;
import java.util.concurrent.*;


/**
 * 定时任务配置类
 */
@EnableScheduling
@Configuration
public class ScheduleConfig implements SchedulingConfigurer {

    private static final String TRACE_ID = TraceId.TRACE_ID;

    @Override
    public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
        // 开启一个线程调度池
        ThreadPoolTaskScheduler taskScheduler = new CustomerThreadPoolTaskScheduler();
        // 设置线程池大小为 1 ：防止多个定时任务并发执行
        taskScheduler.setPoolSize(5);
        taskScheduler.setThreadNamePrefix("codex-schedule-");
        taskScheduler.initialize();
        // taskRegistrar.setTaskScheduler(Executors.newScheduledThreadPool(10));
        taskRegistrar.setScheduler(Executors.newScheduledThreadPool(10));
    }

    public static class CustomerThreadPoolTaskScheduler extends ThreadPoolTaskScheduler {
        @Override
        protected ScheduledExecutorService createExecutor(int poolSize, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler rejectedExecutionHandler) {
            return new CustomerScheduledExecutorService(poolSize, threadFactory, rejectedExecutionHandler);
        }
    }

    public static class CustomerScheduledExecutorService extends ScheduledThreadPoolExecutor {
        @Override
        protected void beforeExecute(Thread t, Runnable r) {
            // 执行前绑定链路追踪ID到线程上下文
            MDC.put(TRACE_ID, UUID.randomUUID().toString().replace("-", ""));
        }

        @Override
        protected void afterExecute(Runnable r, Throwable t) {
            // 执行后清除链路追踪ID
            MDC.remove(TRACE_ID);
        }

        public CustomerScheduledExecutorService(int corePoolSize, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler handler) {
            super(corePoolSize, threadFactory, handler);
        }
    }
}
